//
// Created by mio on 25-5-14.
//

#include "RtmpPublisher.h"

std::shared_ptr<RtmpPublisher> RtmpPublisher::Create(EventLoop *loop) {
    std::shared_ptr<RtmpPublisher> publisher(new RtmpPublisher(loop));
    return publisher;
}

RtmpPublisher::RtmpPublisher(EventLoop *event_loop) : event_loop_(event_loop) {
}


RtmpPublisher::~RtmpPublisher() {
}

int RtmpPublisher::SetMediaInfo(MediaInfo media_info)
{
    media_info_ = media_info;

    if (media_info_.audio_codec_id == RTMP_CODEC_ID_AAC)
    {
        if (media_info_.audio_specific_config_size > 0)
        {
            aac_sequence_header_size_ = media_info_.audio_specific_config_size + 2;
            aac_sequence_header_.reset(new char[aac_sequence_header_size_]);
            uint8_t* data = (uint8_t*)aac_sequence_header_.get();
            data[0] = 0xAF; // 1010 (AAC) 11 (44kHz) 1 (16bit) 1 (stereo)
            data[1] = 0;    // AACPacketType: 0 (sequence header)

            memcpy(data + 2, media_info_.audio_specific_config.get(), media_info_.audio_specific_config_size);
        }
        else { media_info_.audio_codec_id = 0; }
    }

    if (media_info_.video_codec_id == RTMP_CODEC_ID_H264)
    {
        if (media_info_.sps_size > 0 && media_info_.pps_size > 0)
        {
            avc_sequence_header_.reset(new char[4096], std::default_delete<char[]>());
            uint8_t* data = (uint8_t*)avc_sequence_header_.get();
            uint32_t index = 0;

            data[index++] = 0x17; // 1:I frame   7:AVC
            data[index++] = 0x00; // AVC Sequence header

            data[index++] = 0x00;
            data[index++] = 0x00;
            data[index++] = 0x00; // fill in 0

            /* AVCDecoderConfigurationVersion */
            data[index++] = 0x01;                       // ConfigurationVersion
            data[index++] = media_info_.sps.get()[1];   // AVCProfileIndication
            data[index++] = media_info_.sps.get()[2];   // Profile compatibility
            data[index++] = media_info_.sps.get()[3];   // AVCLevel Indication
            data[index++] = 0xFF;                       // LengthSizeMinusOne

            /* SPS nums */
            data[index++] = 0xE1;

            /* SPS data length */
            data[index++] = media_info_.sps_size >> 8;
            data[index++] = media_info_.sps_size & 0xFF;

            /* SPS data */
            memcpy(data + index, media_info_.sps.get(), media_info_.sps_size);
            index += media_info_.sps_size;

            /* PPS nums */
            data[index++] = 0x01;
            data[index++] = media_info_.pps_size >> 8;
            data[index++] = media_info_.pps_size & 0xFF;

            /* PPS data */
            memcpy(data + index, media_info_.pps.get(), media_info_.pps_size);
            index += media_info_.pps_size;

            avc_sequence_header_size_ = index;
        }
    }

    return 0;
}

int RtmpPublisher::OpenUrl(std::string url, int msec)
{
    if (ParseRtmpUrl(url) != 0) return -1;

    if (rtmp_conn_) // 链接存在, 就需要断开连接
    {
        std::shared_ptr<RtmpConnection> conn = rtmp_conn_;
        rtmp_conn_ = nullptr;
        conn->DisConnect();
    }

    // 创建连接
    TcpSocket socket;
    socket.Create();
    if (!socket.Connect(ip_, port_)) {
        socket.Close();
        return -1;
    }

    rtmp_conn_.reset(new RtmpConnection(shared_from_this(), event_loop_->GetTaskScheduler().get(), socket.GetSocket()));

    rtmp_conn_->Handshake();

    std::this_thread::sleep_for(std::chrono::milliseconds(2000));

    return 0;
}

int RtmpPublisher::PushVideoFrame(uint8_t *data, uint32_t size)
{
    if (rtmp_conn_ == nullptr || rtmp_conn_->IsClosed() || size <= 5) return -1;

    if (media_info_.video_codec_id == RTMP_CODEC_ID_H264)
    {
        if (!has_key_frame_)
        {
            if (this->IsKeyFrame(data, size))
            {
                has_key_frame_ = true;
                rtmp_conn_->SendVideoData(0, avc_sequence_header_, avc_sequence_header_size_);
                rtmp_conn_->SendAudioData(0, aac_sequence_header_, aac_sequence_header_size_);
            }
            else return 0;
        }
    }

    uint64_t timestamp = timestamp_.Elapsed();

    std::shared_ptr<char> payload(new char[size + 4096], std::default_delete<char[]>());
    uint32_t payload_size = 0;

    uint8_t* body = (uint8_t*)payload.get();
    uint32_t index = 0;
    body[index++] = this->IsKeyFrame(data, size) ? 0x17 : 0x27; // 1:I frame  2:P frame  7:AVC
    body[index++] = 1;  // AVC NALU

    body[index++] = 0;
    body[index++] = 0;
    body[index++] = 0;

    // NALU size
    body[index++] = (size >> 24) & 0xFF;
    body[index++] = (size >> 16) & 0xFF;
    body[index++] = (size >> 8) & 0xFF;
    body[index++] = size & 0xFF;

    // NALU data
    memcpy(body + index, data, size);
    index += size;
    payload_size = index;
    rtmp_conn_->SendVideoData(timestamp, payload, payload_size);
}

int RtmpPublisher::PushAudioFrame(uint8_t *data, uint32_t size)
{
    if (rtmp_conn_ == nullptr || rtmp_conn_->IsClosed() || size <= 0) return -1;

    if (media_info_.video_codec_id == RTMP_CODEC_ID_H264 && has_key_frame_) {
        uint64_t timestamp = timestamp_.Elapsed();
        uint32_t payload_size = size + 2;
        std::shared_ptr<char> payload(new char[size + 2], std::default_delete<char[]>());

        payload.get()[0] = 0xAF;
        payload.get()[1] = 0x01; // AACPacketType: 1 (raw AAC frame)
        memcpy(payload.get() + 2, data, size);
        rtmp_conn_->SendAudioData(timestamp, payload, payload_size);
    }
    return 0;
}

void RtmpPublisher::Close()
{
    if (rtmp_conn_)
    {
        std::shared_ptr<RtmpConnection> conn = rtmp_conn_;
        rtmp_conn_ = nullptr;
        conn->DisConnect();
        has_key_frame_ = false;
    }
}

bool RtmpPublisher::IsConnected()
{
    if (rtmp_conn_) return !(rtmp_conn_->IsClosed());
    return false;
}


bool RtmpPublisher::IsKeyFrame(uint8_t *data, uint32_t size)
{
    int startcode = 0;
    //判断关键帧 startcode 3 4
    if (data[0] == 0 && data[1] == 0 && data[2] == 0) startcode = 3;
    else if (data[0] == 0 && data[1] == 0 && data[2] == 0 && data[3] == 0) startcode = 4;

    int type = data[startcode] & 0x1F;
    if (type == 0 || type == 7)  return true; // 关键帧
    return false;
}
